#!/usr/bin/env python3
# 未完成代码

import rospy
import math
import tf
from geometry_msgs.msg import Twist
from sensor_msgs.msg import PointCloud2, PointField
from tf.transformations import quaternion_from_euler
from geometry_msgs.msg import PoseWithCovarianceStamped
from tf.transformations import euler_from_quaternion
from sensor_msgs import point_cloud2
from std_msgs.msg import Header


class NavSimulate3D:
    def __init__(self):
        rospy.init_node("nav_simulate_3d", anonymous=False)

        # 初始化位姿（新增 z）
        self.x = rospy.get_param("~initial_x", 0.0)
        self.y = rospy.get_param("~initial_y", 0.0)
        self.z = rospy.get_param("~initial_z", 0.0) 
        self.yaw = rospy.get_param("~initial_yaw", 0.0)

        # 存储速度
        self.linear_velocity = 0.0
        self.angular_velocity = 0.0

        # 全局点云
        self.global_pointcloud = None

        # 上次时间戳
        self.last_time = rospy.Time.now()

        # 订阅速度指令
        rospy.Subscriber("/cmd_vel", Twist, self.cmd_vel_callback)

        # 订阅3D点云地图
        rospy.Subscriber("/map3d", PointCloud2, self.pointcloud_callback)

        # 发布局部点云
        self.local_pointcloud_pub = rospy.Publisher(
            "/pc_accumulate", PointCloud2, queue_size=10
        )

        # 订阅初始位姿
        rospy.Subscriber(
            "/initialpose", PoseWithCovarianceStamped, self.pose_reset_callback
        )

        # TF 广播器
        self.br = tf.TransformBroadcaster()

        # 每 10Hz 执行一次
        self.timer = rospy.Timer(rospy.Duration(1.0 / 10.0), self.timer_callback)

    def cmd_vel_callback(self, msg):
        self.linear_velocity = msg.linear.x
        self.angular_velocity = msg.angular.z

    def pointcloud_callback(self, msg):
        self.global_pointcloud = msg

    def timer_callback(self, event):
        current_time = rospy.Time.now()
        dt = (current_time - self.last_time).to_sec()
        if dt > 0:
            # 更新 x, y, yaw
            self.x += self.linear_velocity * math.cos(self.yaw) * dt
            self.y += self.linear_velocity * math.sin(self.yaw) * dt
            self.yaw += self.angular_velocity * dt

            # 垂直控制逻辑
            if self.global_pointcloud:
                self.update_z_based_on_pointcloud()

        self.last_time = current_time

        # 发布 TF 变换
        self.publish_transform()

        # 发布局部点云
        if self.global_pointcloud:
            local_pointcloud = self.create_local_pointcloud()
            self.local_pointcloud_pub.publish(local_pointcloud)

    def update_z_based_on_pointcloud(self):
        # 定义搜索范围（单位：米）
        search_radius = 1.0
        min_z = -100.0

        # 在点云中查找当前 x, y 附近的点
        points = []
        try:
            for p in point_cloud2.read_points(
                self.global_pointcloud, field_names=("x", "y", "z"), skip_nans=True
            ):
                if (
                    abs(p[0] - self.x) < search_radius
                    and abs(p[1] - self.y) < search_radius
                ):
                    points.append(p)
        except Exception as e:
            rospy.logwarn(f"Error reading points: {e}")
            return

        if points:
            # 有地面点，抬升 z 到平均高度
            avg_z = sum(p[2] for p in points) / len(points)
            self.z = max(avg_z, self.z + 0.01)  # 缓慢上升
        else:
            # 无地面点，下坠
            self.z -= 0.1
            if self.z < min_z:
                self.reset_to_origin()

    def reset_to_origin(self):
        self.x = 0.0
        self.y = 0.0
        self.z = 0.0
        self.yaw = 0.0
        rospy.loginfo("Reset to origin due to falling below -100m.")

    def publish_transform(self):
        q = quaternion_from_euler(0, 0, self.yaw)
        self.br.sendTransform(
            (self.x, self.y, self.z), q, rospy.Time.now(), "body", "map"
        )

    def create_local_pointcloud(self):
        header = Header(stamp=rospy.Time.now(), frame_id="body")
        points = []

        # 定义局部区域范围（例如 10m × 10m × 10m）
        x_range = (-5.0, 5.0)
        y_range = (-5.0, 5.0)
        z_range = (-5.0, 5.0)

        # 计算当前姿态的余弦和正弦值
        cos_yaw = math.cos(self.yaw)
        sin_yaw = math.sin(self.yaw)

        try:
            for p in point_cloud2.read_points(
                self.global_pointcloud, field_names=("x", "y", "z"), skip_nans=True
            ):
                # 计算全局点相对于机器人的坐标
                dx = p[0] - self.x
                dy = p[1] - self.y
                dz = p[2] - self.z

                # 应用旋转变换到局部坐标系
                x_local = dx * cos_yaw + dy * sin_yaw
                y_local = -dx * sin_yaw + dy * cos_yaw
                z_local = dz  # z轴仅平移

                # 检查是否在局部范围内
                if (
                    x_range[0] < x_local < x_range[1]
                    and y_range[0] < y_local < y_range[1]
                    and z_range[0] < z_local < z_range[1]
                ):
                    points.append([x_local, y_local, z_local])
        except Exception as e:
            rospy.logwarn(f"Error creating local cloud: {e}")

        return point_cloud2.create_cloud_xyz32(header, points)

    def pose_reset_callback(self, msg):
        self.x = msg.pose.pose.position.x
        self.y = msg.pose.pose.position.y
        self.z = msg.pose.pose.position.z  # 新增 z

        orientation_q = msg.pose.pose.orientation
        orientation_list = [
            orientation_q.x,
            orientation_q.y,
            orientation_q.z,
            orientation_q.w,
        ]
        _, _, self.yaw = euler_from_quaternion(orientation_list)

        self.last_time = rospy.Time.now()


if __name__ == "__main__":
    try:
        ns = NavSimulate3D()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
